import argparse
import pickle

import torch
import torch_npu
from loguru import logger


def run_pt_api(api, input_path, output_path):
    try:
        data = pickle.load(open(input_path, 'rb'))

        cpu_data = data
        for key in cpu_data:
            if torch.is_tensor(cpu_data[key]):
                cpu_data[key] = cpu_data[key].to('cpu')
        cpu_result = eval(api)(**cpu_data)

        gpu_data = data
        for key in gpu_data:
            if torch.is_tensor(gpu_data[key]):
                gpu_data[key] = gpu_data[key].to('cuda')
        gpu_result = eval(api)(**gpu_data)

        npu_data = data
        for key in npu_data:
            if torch.is_tensor(npu_data[key]):
                npu_data[key] = npu_data[key].to('npu')
        npu_result = eval(api)(**npu_data)

        data['cpu_result'] = cpu_result.to('cpu')
        data['gpu_result'] = gpu_result.to('cpu')
        data['npu_result'] = npu_result.to('cpu')
        pickle.dump(data, open(output_path, 'wb'))
    except Exception as e:
        logger.error(e)


if __name__ == '__main__':
    # option = {'ACL_PRECISION_MODE': 'must_keep_origin_dtype'}
    # torch.npu.set_option(option)
    parser = argparse.ArgumentParser()
    parser.add_argument('--api', type=str)
    parser.add_argument('--input_path', type=str)
    parser.add_argument('--output_path', type=str)
    args = parser.parse_args()
    run_pt_api(args.api, args.input_path, args.output_path)
